package com.apex.spark.batch;

import com.apex.env.Execution;
import com.apex.spark.SparkEnvironment;
import com.apex.spark.transform.SparkTransform;
import com.typesafe.config.Config;
import org.apache.commons.lang3.StringUtils;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * spark 批处理执行入口
 */
public class SparkBatchExecution implements Execution<SparkBatchSource, SparkTransform, SparkBatchSink> {
    //上下文执行环境
    private final SparkEnvironment sparkEnvironment;
    //配置文件管理
    private Config config;
    //程序执行状态
    private boolean status ;

    /**
     * 配置执行环境
     * @param sparkEnvironment 上下文执行环境
     */
    public SparkBatchExecution(SparkEnvironment sparkEnvironment) {
        this.sparkEnvironment = sparkEnvironment;
    }

    /**
     * 数据处理方法
     * @param sources 数据源
     * @param transforms 业务处理
     * @param sinks 写入目标数据
     */
    @Override
    public void start(List<SparkBatchSource> sources, List<SparkTransform> transforms, List<SparkBatchSink> sinks) {
        try{
            //用户自定义spark查询临时表名称
            String tableName = sparkEnvironment.getQueryTable();
            //保存源端数据输入集
            List<Dataset<Row>> data = new ArrayList<>();
            //数据处理结果
            Dataset<Row> result = null;
            //判断source是否存在
            if (!sources.isEmpty()){
                //获取数据
                for (SparkBatchSource source : sources) {
                    Dataset<Row> dataset = source.getData(sparkEnvironment);
                    data.add(dataset);
                    //todo 注册数据集为spark临时表
                    if (!StringUtils.isBlank(tableName)){
                        registerTempView(tableName, dataset);
                    }
                }
                //获取数据源,并发送给下游算子
                Dataset<Row> input = data.get(0);

                //数据处理方法算子transform阶段
                if (!transforms.isEmpty()){
                    for (SparkTransform transform : transforms) {
                        result = transform.process(sparkEnvironment, input);
                    }
                }

                //对上游数据进行入库操作,需要判断上游transform阶段发来的数据集是否为空,如果为空则表示入库操作无需处理
                for (SparkBatchSink sink : sinks) {
                    //对transform阶段发来的数据集进行判断
                    if (Objects.isNull(result)){
                        sink.outputBatch(sparkEnvironment, input);
                    }else {
                        sink.outputBatch(sparkEnvironment,result);
                    }

                }
            }

        }catch (Exception e){
            //todo
            status = false;
            e.printStackTrace();
        }
    }

    /**
     * 程序执行状态
     * @return
     */
    @Override
    public boolean status() {
        return status;
    }

    /**
     *
     * @param plugin 执行环境
     */
    @Override
    public void prepare(Void plugin) {

    }

    /**
     *
     * @return config
     */
    @Override
    public Config getConfig() {
        return config;
    }

    /**
     *
     * @param config config
     */
    @Override
    public void setConfig(Config config) {
        this.config = config;
    }

    /**
     * 注册临时表结果集
     * @param tableName 表名称
     * @param dataSet 查询结果集
     */
    private void registerTempView(String tableName, Dataset<Row> dataSet) {
        dataSet.createOrReplaceTempView(tableName);
    }
}
